from uav_2d import Uav2DEnv
from uav_2d.wrappers.raster_wrapper import RasterWrapper


def rl_test(env: Uav2DEnv, episodes: int = 3, render: bool = False):
    for i in range(episodes):
        obs, _ = env.reset()
        done = False

        while not done:
            action = env.action_space.sample()
            obs, reward, terminate, truncate, info = env.step(action)
            done = terminate["__all__"] or truncate["__all__"]
            if render:
                env.render()
        else:
            print(f"Episode {i} done")
    env.close()


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--render", action="store_true", default=True)
    parser.add_argument("--add_raster", action="store_true")
    parser.add_argument("--width", type=int, default=1000)
    parser.add_argument("--height", type=int, default=1000)
    parser.add_argument("--episodes", type=int, default=15)
    parser.add_argument("--num_agents", type=int, default=3)
    parser.add_argument("--fixed_obstacle_number", type=int, default=20)
    parser.add_argument("--occur_fixed_number", type=int, default=1)
    parser.add_argument("--occur_number_max", type=int, default=3)
    parser.add_argument("--dynamic_obstacle_number", type=int, default=20)
    args = parser.parse_args()

    dim = (args.width, args.height)
    env = Uav2DEnv(
        {
            "num_agents": args.num_agents,
            "dimensions": dim,
            "fixed_obstacles": args.fixed_obstacle_number,
            "dynamic_obstacles": args.dynamic_obstacle_number,
            "occur_obstacles": args.occur_fixed_number,
            "occur_number_max": args.occur_number_max,
            "prevent_stiff": False,
            "show_windows": args.render,
            "use_lidar": True,
            "draw_lidar": True,
            "lidar_range": 250,
            "lidar_rays": 21,
            "field_of_view": 210,
            "center_obstacles": False,
        },
    )
    if args.add_raster:
        env = RasterWrapper(env)

    rl_test(env, args.episodes, args.render)
